
/****************************************************************************/
/*Copyright (c) 2011, Florent DEVILLE.                                      */
/*All rights reserved.                                                      */
/*                                                                          */
/*Redistribution and use in source and binary forms, with or without        */
/*modification, are permitted provided that the following conditions        */
/*are met:                                                                  */
/*                                                                          */
/* - Redistributions of source code must retain the above copyright         */
/*notice, this list of conditions and the following disclaimer.             */
/* - Redistributions in binary form must reproduce the above                */
/*copyright notice, this list of conditions and the following               */
/*disclaimer in the documentation and/or other materials provided           */
/*with the distribution.                                                    */
/* - The names of its contributors cannot be used to endorse or promote     */
/*products derived from this software without specific prior written        */
/*permission.                                                               */
/* - The source code cannot be used for commercial purposes without         */ 
/*its contributors' permission.                                             */
/*                                                                          */
/*THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS       */
/*"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT         */
/*LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS         */
/*FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE            */
/*COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,       */
/*INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,      */
/*BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;          */
/*LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER          */
/*CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT        */
/*LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN         */
/*ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE           */
/*POSSIBILITY OF SUCH DAMAGE.                                               */
/****************************************************************************/

#define WIN32_LEAN_AND_MEAN		// Exclude rarely-used stuff from Windows headers
// Windows Header Files:
// C RunTime Header Files
#include <stdlib.h>
#include <malloc.h>
#include <memory.h>
#include <tchar.h>
#include <windows.h>
#include <stdio.h>

#include "RTAviFile.h"


#ifndef __countof
#define __countof(x)	((sizeof(x)/sizeof(x[0])))
#endif

namespace RT
{

	RTAviFile:: RTAviFile(LPCTSTR lpszFileName /* =_T("Output.avi") */, 
				DWORD dwCodec /* = mmioFOURCC('M','P','G','4') */,
				DWORD dwFrameRate /* = 1 */)
	{

		AVIFileInit();

		m_hHeap=NULL;
		m_hAviDC=NULL;
		m_lpBits=NULL;
		m_lSample=NULL;
		m_pAviFile=NULL;
		m_pAviStream=NULL;
		m_pAviCompressedStream=NULL;

		m_dwFCCHandler = dwCodec;
		m_dwFrameRate = dwFrameRate;

		//_tcscpy(m_szFileName, lpszFileName);
		wcscpy_s(m_szFileName, lpszFileName);

		//_tcscpy(m_szErrMsg, _T("Method Succeeded"));
		wcscpy_s(m_szErrMsg, _T("Method Succeeded"));

		m_szErrMsg[__countof(m_szErrMsg)-1] = _T('\0');

		pAppendFrame[0]= &RTAviFile::AppendDummy;			// VC8 requires & for Function PoI32er; Remove it if your compiler complains;
		pAppendFrame[1]= &RTAviFile::AppendFrameFirstTime;
		pAppendFrame[2]= &RTAviFile::AppendFrameUsual;

		pAppendFrameBits[0]=&RTAviFile::AppendDummy;
		pAppendFrameBits[1]=&RTAviFile::AppendFrameFirstTime;
		pAppendFrameBits[2]=&RTAviFile::AppendFrameUsual;

		m_nAppendFuncSelector=1;		//0=Dummy	1=FirstTime	2=Usual
	}

	RTAviFile::~RTAviFile(void)
	{
		ReleaseMemory();

		AVIFileExit();
	}

	void RTAviFile::ReleaseMemory()
	{
		m_nAppendFuncSelector=0;		//PoI32 to DummyFunction

		if(m_hAviDC)
		{
			DeleteDC(m_hAviDC);
			m_hAviDC=NULL;
		}
		if(m_pAviCompressedStream)
		{
			//AVIStreamRelease(m_pAviCompressedStream);
			AVIStreamClose(m_pAviCompressedStream);
			m_pAviCompressedStream=NULL;
		}
		if(m_pAviStream)
		{
			AVIStreamRelease(m_pAviStream);
			m_pAviStream=NULL;
		}
		if(m_pAviFile)
		{
			AVIFileRelease(m_pAviFile);
			m_pAviFile=NULL;
		}
		if(m_lpBits)
		{
			HeapFree(m_hHeap,HEAP_NO_SERIALIZE,m_lpBits);
			m_lpBits=NULL;
		}
		if(m_hHeap)
		{
			HeapDestroy(m_hHeap);
			m_hHeap=NULL;
		}
	}

	void RTAviFile::SetErrorMessage(LPCTSTR lpszErrorMessage)
	{
		//_tcsncpy(m_szErrMsg, lpszErrorMessage, __countof(m_szErrMsg)-1);
		wcsncpy_s(m_szErrMsg, lpszErrorMessage, __countof(m_szErrMsg)-1);
	}
	
	HRESULT RTAviFile::InitMovieCreation(I32 nFrameWidth, I32 nFrameHeight, I32 nBitsPerPixel)
	{
		I32	nMaxWidth=GetSystemMetrics(SM_CXSCREEN),nMaxHeight=GetSystemMetrics(SM_CYSCREEN);

		m_hAviDC = CreateCompatibleDC(NULL);
		if(m_hAviDC==NULL)	
		{
			SetErrorMessage(_T("Unable to Create Compatible DC"));
			return E_FAIL;
		}
	
		if(nFrameWidth > nMaxWidth)	nMaxWidth= nFrameWidth;
		if(nFrameHeight > nMaxHeight)	nMaxHeight = nFrameHeight;

		m_hHeap=HeapCreate(HEAP_NO_SERIALIZE, nMaxWidth*nMaxHeight*4, 0);
		if(m_hHeap==NULL)
		{
			SetErrorMessage(_T("Unable to Create Heap"));
			return E_FAIL;
		}
	
		m_lpBits=HeapAlloc(m_hHeap, HEAP_ZERO_MEMORY|HEAP_NO_SERIALIZE, nMaxWidth*nMaxHeight*4);
		if(m_lpBits==NULL)	
		{	
			SetErrorMessage(_T("Unable to Allocate Memory on Heap"));
			return E_FAIL;
		}

		if(FAILED(AVIFileOpen(&m_pAviFile, m_szFileName, OF_CREATE|OF_WRITE, NULL)))
		{
			SetErrorMessage(_T("Unable to Create the Movie File"));
			return E_FAIL;
		}

		ZeroMemory(&m_AviStreamInfo,sizeof(AVISTREAMINFO));
		m_AviStreamInfo.fccType		= streamtypeVIDEO;
		m_AviStreamInfo.fccHandler	= m_dwFCCHandler;
		m_AviStreamInfo.dwScale		= 1;
		m_AviStreamInfo.dwRate		= m_dwFrameRate;	// Frames Per Second;
		m_AviStreamInfo.dwQuality	= (DWORD)-1;				// Default Quality
		m_AviStreamInfo.dwSuggestedBufferSize = nMaxWidth*nMaxHeight*4;
		SetRect(&m_AviStreamInfo.rcFrame, 0, 0, nFrameWidth, nFrameHeight);
		//_tcscpy(m_AviStreamInfo.szName, _T("Video Stream"));
		wcscpy_s(m_AviStreamInfo.szName, _T("Video Stream"));

		if(FAILED(AVIFileCreateStream(m_pAviFile,&m_pAviStream,&m_AviStreamInfo)))
		{
			SetErrorMessage(_T("Unable to Create Video Stream in the Movie File"));
			return E_FAIL;
		}

		ZeroMemory(&m_AviCompressOptions,sizeof(AVICOMPRESSOPTIONS));
		m_AviCompressOptions.fccType=streamtypeVIDEO;
		m_AviCompressOptions.fccHandler=m_AviStreamInfo.fccHandler;
		m_AviCompressOptions.dwFlags=AVICOMPRESSF_KEYFRAMES|AVICOMPRESSF_VALID;//|AVICOMPRESSF_DATARATE;
		m_AviCompressOptions.dwKeyFrameEvery=1;
		//m_AviCompressOptions.dwBytesPerSecond=1000/8;
		//m_AviCompressOptions.dwQuality=100;

		HRESULT hr = AVIMakeCompressedStream(&m_pAviCompressedStream,m_pAviStream,&m_AviCompressOptions,NULL);
		if( hr!= AVIERR_OK )
		{
			// One reason this error might occur is if you are using a Codec that is not 
			// available on your system. Check the mmioFOURCC() code you are using and make
			// sure you have that codec installed properly on your machine.
			SetErrorMessage(_T("Unable to Create Compressed Stream: Check your CODEC options"));
			return E_FAIL;
		}

		BITMAPINFO bmpInfo;
		ZeroMemory(&bmpInfo,sizeof(BITMAPINFO));
		bmpInfo.bmiHeader.biPlanes		= 1;
		bmpInfo.bmiHeader.biWidth		= nFrameWidth;
		bmpInfo.bmiHeader.biHeight		= nFrameHeight;
		bmpInfo.bmiHeader.biCompression	= BI_RGB;
		bmpInfo.bmiHeader.biBitCount	= (WORD)nBitsPerPixel;
		bmpInfo.bmiHeader.biSize		= sizeof(BITMAPINFOHEADER);
		bmpInfo.bmiHeader.biSizeImage	= bmpInfo.bmiHeader.biWidth*bmpInfo.bmiHeader.biHeight*bmpInfo.bmiHeader.biBitCount/8;

		if(FAILED(AVIStreamSetFormat(m_pAviCompressedStream,0,(LPVOID)&bmpInfo, bmpInfo.bmiHeader.biSize)))
		{
			// One reason this error might occur is if your bitmap does not meet the Codec requirements.
			// For example, 
			//   your bitmap is 32bpp while the Codec supports only 16 or 24 bpp; Or
			//   your bitmap is 274 * 258 size, while the Codec supports only sizes that are powers of 2; etc...
			// Possible solution to avoid this is: make your bitmap suit the codec requirements,
			// or Choose a codec that is suitable for your bitmap.
			SetErrorMessage(_T("Unable to Set Video Stream Format"));
			return E_FAIL;
		}

		return S_OK;	// Everything went Fine
	}


	HRESULT	RTAviFile::AppendFrameFirstTime(HBITMAP	hBitmap)
	{
		BITMAP Bitmap;

		GetObject(hBitmap, sizeof(BITMAP), &Bitmap);

		if(SUCCEEDED(InitMovieCreation( Bitmap.bmWidth, 
										Bitmap.bmHeight, 
										Bitmap.bmBitsPixel)))
		{
			m_nAppendFuncSelector=2;		//PoI32 to the UsualAppend Function

			return AppendFrameUsual(hBitmap);
		}

		ReleaseMemory();

		return E_FAIL;
	}

	HRESULT RTAviFile::AppendFrameUsual(HBITMAP hBitmap)
	{
		BITMAPINFO	bmpInfo;

		bmpInfo.bmiHeader.biBitCount=0;
		bmpInfo.bmiHeader.biSize=sizeof(BITMAPINFOHEADER);
	
		GetDIBits(m_hAviDC,hBitmap,0,0,NULL,&bmpInfo,DIB_RGB_COLORS);

		bmpInfo.bmiHeader.biCompression=BI_RGB;	

		GetDIBits(m_hAviDC,hBitmap,0,bmpInfo.bmiHeader.biHeight,m_lpBits,&bmpInfo,DIB_RGB_COLORS);

		if(FAILED(AVIStreamWrite(m_pAviCompressedStream,m_lSample++,1,m_lpBits,bmpInfo.bmiHeader.biSizeImage,0,NULL,NULL)))
		{
			SetErrorMessage(_T("Unable to Write Video Stream to the output Movie File"));

			ReleaseMemory();

			return E_FAIL;
		}

		return S_OK;
	}

	HRESULT RTAviFile::AppendDummy(HBITMAP)
	{
		return E_FAIL;
	}

	HRESULT RTAviFile::AppendNewFrame(HBITMAP hBitmap)
	{
		return (this->*pAppendFrame[m_nAppendFuncSelector])((HBITMAP)hBitmap);
	}

	HRESULT	RTAviFile::AppendNewFrame(I32 nWidth, I32 nHeight, LPVOID pBits,I32 nBitsPerPixel)
	{
		return (this->*pAppendFrameBits[m_nAppendFuncSelector])(nWidth,nHeight,pBits,nBitsPerPixel);
	}

	HRESULT	RTAviFile::AppendFrameFirstTime(I32 nWidth, I32 nHeight, LPVOID pBits,I32 nBitsPerPixel)
	{
		if(SUCCEEDED(InitMovieCreation(nWidth, nHeight, nBitsPerPixel)))
		{
			m_nAppendFuncSelector=2;		//PoI32 to the UsualAppend Function

			return AppendFrameUsual(nWidth, nHeight, pBits, nBitsPerPixel);
		}

		ReleaseMemory();

		return E_FAIL;
	}

	HRESULT	RTAviFile::AppendFrameUsual(I32 nWidth, I32 nHeight, LPVOID pBits,I32 nBitsPerPixel)
	{
		DWORD	dwSize=nWidth*nHeight*nBitsPerPixel/8;

		if(FAILED(AVIStreamWrite(m_pAviCompressedStream,m_lSample++,1,pBits,dwSize,0,NULL,NULL)))
		{
			SetErrorMessage(_T("Unable to Write Video Stream to the output Movie File"));

			ReleaseMemory();

			return E_FAIL;
		}

		return S_OK;
	}

	HRESULT	RTAviFile::AppendDummy(I32 /*nWidth*/, I32 /*nHeight*/, LPVOID /*pBits*/,I32 /*nBitsPerPixel*/)
	{
		return E_FAIL;
	}

}